package edu.mit.compilers.assembly.x64;

import edu.mit.compilers.assembly.x64.basic.*;
import edu.mit.compilers.lowercode.FieldInfo;
import edu.mit.compilers.lowercode.MethodInfo;
import edu.mit.compilers.lowercode.SymbolTable;
import edu.mit.compilers.lowercode.code.ThreeAddressCodeList;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;

public class X64EntryAllocator {
    public final static long VAR_ALIGN_SIZE = 8;
    public final static long RSP_ALIGN_SIZE = 16;   // when calling method suck as "printf", rsp must be 16bytes-aligned

    public static class AllocateResult {
        private boolean ok;
        private X64Entry entry;

        public boolean ok() {
            return ok;
        }

        public X64Entry getEntry() {
            assert ok;
            return ok ? entry : null;
        }
    }

    private final String methodName;
    private final SymbolTable symbolTable;
    private final ThreeAddressCodeList methodCodes;
    
    private X64AllocateMap allocatedMap = null;
    private Set<String> freeTempRegisters = null;
    
    public X64EntryAllocator(String methodName, SymbolTable symbolTable, ThreeAddressCodeList methodCodes) {
        this.methodName = methodName;
        this.symbolTable = symbolTable;
        this.methodCodes = methodCodes;
        Init();
    }

    public long getStackFrameSize() {
        return allocatedMap.getAllocatedStackSize();
    }

    private void Init() {
        allocatedMap = staticAllocate(methodName, symbolTable, methodCodes);
        freeTempRegisters = new HashSet<>();
        freeTempRegisters.add("rax");
        freeTempRegisters.add("rcx");
        freeTempRegisters.add("rdx");
        freeTempRegisters.add("rsi");
        freeTempRegisters.add("rdi");
        freeTempRegisters.add("r8");
        freeTempRegisters.add("r9");
        freeTempRegisters.add("r10");
        freeTempRegisters.add("r11");
    }

    // allocate a address for variable @variable
    public AllocateResult allocate(String variable) {
        AllocateResult ret = new AllocateResult();
        ret.entry = allocatedMap.get(variable);
        ret.ok = ret.entry != null;
        assert ret.ok : "variable " + variable + " not exist in allocate map";
        return ret;
    }

    public AllocateResult allocateGlobal(String variable, int size) {
        assert size == 8;
        AllocateResult ret = new AllocateResult();
        ret.entry = allocateGlobalEntry(variable, size);
        ret.ok = true;
        return ret;
    }

    public AllocateResult allocateScaled(String baseRegister, String offsetRegister, int scaled) {
        assert scaled == 8;
        AllocateResult ret = new AllocateResult();
        ret.entry = allocateScaledEntry(baseRegister, offsetRegister, scaled);
        ret.ok = true;
        return ret;
    }

    public void releaseTempRegister(X64Register register) {
        // we only consider r9-r11
        if (canUsedAsTemp(register)) {
            freeTempRegisters.add(register.name);
        }
    }

    private static boolean canUsedAsTemp(X64Register register) {
        String name = register.name;
        return register.temporary && name.charAt(0) == 'r' && !name.equals("rbp") && !name.equals("rip");
    }

    // allocate temp register, we assume this register will be released after moved
    public AllocateResult allocateTempRegister(int size) {
        assert size == 8;
        AllocateResult ret = new AllocateResult();
        ret.entry = allocateTempRegisterEntry();
        ret.ok = ret.entry != null;
        return ret;
    }

    public AllocateResult allocateRunStackPointer() {
        AllocateResult ret = new AllocateResult();
        ret.entry = allocateRegisterEntry("rsp");
        ret.ok = true;
        return ret;
    }

    public AllocateResult allocatePopArgInStack(int i) {
        AllocateResult ret = new AllocateResult();
        ret.entry = allocateStackEntry(8, i * 8L + 16);
        ret.ok = true;
        return ret;
    }

    private X64Entry allocateTempRegisterEntry() {
        if (freeTempRegisters.isEmpty()) {
            return null;
        }
        return allocateRegisterEntry(freeTempRegisters.iterator().next());
    }

    public AllocateResult allocateRetRegister(int size) {
        assert size == 8;
        AllocateResult ret = new AllocateResult();
        ret.entry = allocateRegisterEntry("rax");
        ret.ok = true;
        return ret;
    }

    public ArrayList<AllocateResult> allocateRegistersForDiv() {
        ArrayList<AllocateResult> ret = new ArrayList<>();
        AllocateResult res1 = new AllocateResult();
        AllocateResult res2 = new AllocateResult();
        res1.entry = allocateRegisterEntry("rax");
        res2.entry = allocateRegisterEntry("rdx");
        res1.ok = res2.ok = true;
        ret.add(res1);
        ret.add(res2);
        return ret;
    }

    public AllocateResult allocateArgRegister(int pushArgIndex, int size) {
        assert size == 8;
        assert pushArgIndex >= 0 && pushArgIndex < 6;

        AllocateResult ret = new AllocateResult();
        X64Entry entry = allocateRegisterEntry(getKthPushArgsRegisterName(pushArgIndex, size));

        ret.entry = entry;
        ret.ok = true;
        return ret;
    }

    public AllocateResult allocateImmediate(String immediate) {
        AllocateResult ret = new AllocateResult();
        ret.entry = allocateImmEntry(immediate);
        ret.ok = true;
        return ret;
    }

    private X64Entry allocateImmEntry(String immediate) {
        X64Entry entry = new X64Entry();
        entry.location = new X64ImmOperand(immediate);
        entry.size = 8;
        return entry;
    }

    private static X64AllocateMap staticAllocate(String methodName, SymbolTable symbolTable, ThreeAddressCodeList methodCodes) {
        // use trivial allocate strategy
        X64AllocateMap ret = new X64AllocateMap();
        MethodInfo methodInfo = symbolTable.getMethodInfo(methodName);

        long offset = 0;
        for (FieldInfo fieldInfo : methodInfo.getFieldsInfo().values()) {
            long size = fieldInfo.getSize();
            offset -= size;
            offset = -alignUpOffset(-offset, VAR_ALIGN_SIZE);
            X64Entry entry = allocateStackEntry(size, offset);
            ret.put(fieldInfo.getName(), entry);
        }
        offset = -alignUpOffset(-offset, RSP_ALIGN_SIZE);
        ret.setStackAllocatedSize(-offset);
        return ret;
    }

    private static long alignUpOffset(long l, long alignedSize) {
        return (l + alignedSize - 1) / alignedSize * alignedSize;
    }

    private static X64Entry allocateStackEntry(long size, long offset) {
        X64Entry entry = new X64Entry();
        entry.size = size;
        entry.location = new X64RegisterMemOperand(X64Register.RBP, String.valueOf(offset));
        return entry;
    }

    private X64Entry allocateRegisterEntry(String register) {
        X64Register x64Register = X64Register.getRegister(register);
        if (canUsedAsTemp(x64Register)) {
            freeTempRegisters.remove(x64Register.name);
        }
        X64Entry entry = new X64Entry();
        entry.location = new X64RegisterOperand(x64Register);
        entry.size = x64Register.bytes;
        return entry;
    }

    private X64Entry allocateGlobalEntry(String variable, int size) {
        X64Entry entry = new X64Entry();
        entry.location = new X64RegisterMemOperand(X64Register.RIP, variable);
        entry.size = size;
        return entry;
    }

    private X64Entry allocateScaledEntry(String baseRegister, String offsetRegister, int scaled) {
        X64Entry entry = new X64Entry();
        entry.location = new X64ScaledMemOperand(
                X64Register.getRegister(baseRegister),
                X64Register.getRegister(offsetRegister),
                scaled);
        entry.size = scaled;
        return entry;
    }

    private static String getKthPushArgsRegisterName(int kth, int size) {
        // %rdi %rsi %rdx %rcx %r8 %r9
        assert size == 8;
        assert kth < 6 : "not support kth " + kth;
        switch (kth) {
            case 0: return "rdi";
            case 1: return "rsi";
            case 2: return "rdx";
            case 3: return "rcx";
            case 4: return "r8";
            case 5: return "r9";
            default:    return null;
        }
    }
}
